Reuse VJPs in GeomSum and ManualIter fixed-point gradients - #486

actually looking at this a bit more, I wonder if we can also speed up the solvers by creating the following 3 functions to compute vjp's for:

f(state, env) = gauge_fixed_iteration(state, env) # both
g = Base.Fix1(f, state)
h = Base.Fix2(f, env)

It might be worth investigating if these vjp's are faster.

In any case, good catch, this does seem like it should speed up the gradients!

left some minor comments about the docstrings but otherwise good to go for me, definitely also okay to put the other changes in a different PR.

Up to a few very minor naming comments, looks perfect to me.

On the Fix1/Fix2 idea: I don’t think that will make a difference. Zygote differentiates through captured fields of a closure, so the state captured in Fix1(f, state) is still tracked, and TensorOperations will still compute (and force) the ket/bra cotangents on every pullback call. I first tried this approach in #430 hoping for a speedup, but it turned out that capturing the state in a closure gave identical allocations to passing it as an argument, which I was very annoyed about.

The only thing that helped was reading the local sandwich through a custom constant_site getindex helper whose gradient is explicitly ignored, so that the state cotangents are never tracked in the pure-environment VJP used in the linear solves.

I then gave up on doing the same for FixedPointGradient, since plumbing a similar hook all the way deep inside a CTMRG iteration was a total pain. Maybe with Enzyme the Fix1/Fix2 idea would work, but with our current setup it makes no difference.